from typing import Optional


class Trie:  # 字典树 前缀树
    class _Node:
        def __init__(self, is_word: Optional[bool] = False):
            self.is_word = is_word  # "pan" "panda"
            self.next = {}  # char -> node

        def __repr__(self):
            return f'_Node(is_word={self.is_word}, next={self.next!r})'

    def __init__(self):
        self.root = self._Node()  # 根节点不包含任何字符
        self.size = 0

    def __len__(self) -> int:
        return self.size

    def add(self, word: str) -> None:
        cur = self.root  # root doesn't store character
        for char in word:
            if cur.next.get(char) is None:
                cur.next[char] = self._Node()  # Node不保存字符
            cur = cur.next.get(char)  # 这里cur一定不为None

        # 该单词还不存在，如panda存在但是pan不存在，此时如果加入pan,cur.is_word为False,这时应该设置为True
        if not cur.is_word:
            cur.is_word = True
            self.size += 1

    def __contains__(self, word) -> bool:
        cur = self.root
        for char in word:
            if cur.next.get(char) is None:  # next为空字典
                return False
            cur = cur.next.get(char)

        return cur.is_word  # 可能包含完整的字符串，但不一定是单词

    def is_prefix(self, prefix) -> bool:
        cur = self.root
        for char in prefix:
            if cur.next.get(char) is None:  # next为空字典
                return False
            cur = cur.next.get(char)

        return True  # 不需要是单词

    def __str__(self):
        return str(self.root)


if __name__ == '__main__':
    trie = Trie()
    '''
    trie.add('pan')
    print(trie)

    print('pan' in trie)
    print('pa' in trie, trie.is_prefix('pa'))
    print('panda' in trie, trie.is_prefix('panda'))


    trie.add('panda')
    print(trie)
    
    '''
    trie.add('pan')
    trie.add('panda')
    trie.add('door')
    trie.add('deerr')
    print(trie)  # {'d': _Node(next={'o': _Node, 'e': _Node})

    print("pan" in trie) # True
    print("pana" in trie) # False
    print("deer" in trie) # False
    print("deerr" in trie) # True
